package com.study.associationrule

import org.apache.spark.ml.fpm.FPGrowth
import org.apache.spark.sql.SparkSession

/**
 * 关联规则
 *
 * @author stephen
 * @date 2019-08-30 10:04
 */
object FPGrowthDemo {

  def main(args: Array[String]): Unit = {

    val spark = SparkSession.builder()
      .appName(this.getClass.getSimpleName)
      .master("local[*]")
      .getOrCreate()

    spark.sparkContext.setLogLevel("warn")
    import spark.implicits._
    val dataset = spark.createDataset(Seq(
      "1 2 5",
      "1 2 3 5",
      "1 2")
    ).map(t => t.split(" ")).toDF("items")

    val fpgrowth = new FPGrowth().setItemsCol("items").setMinSupport(0.5).setMinConfidence(0.6)
    val model = fpgrowth.fit(dataset)

    // 模型产出结果
    model.freqItemsets.show()

    // 模型生成的关联规则
    model.associationRules.show()

    // 通过关联规则进行预测
    model.transform(dataset).show()
  }
}
